Skip to content

Refactor: 3.1 phase of RESTRUCTURE.md, focussing on setup.sh and related Python dependency & Dockerfile files #2541

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 1 commit into from
Nov 3, 2025
Merged

Refactor: 3.1 phase of RESTRUCTURE.md, focussing on setup.sh and related Python dependency & Dockerfile files #2541

merged 1 commit into from
Nov 3, 2025

Conversation

@SamuelMarks
Copy link
Collaborator

@SamuelMarks SamuelMarks commented Oct 24, 2025 •
edited
Loading

Description

Refactor: 3.1 phase of RESTRUCTURE.md, focussing on setup.sh and related Python dependency & Dockerfile files

Note: merge this before #2361; after which that PR is roughly halved.

Tests

CI and manual:

$ bash ./dependencies/scripts/docker_build_dependency_image.sh DEVICE='tpu' MODE='nightly'
$ bash ./dependencies/scripts/docker_build_dependency_image.sh DEVICE='tpu' MODE='stable'
$ export MODEL_NAME='llama3_1_70b_8192_synthetic' \
         PROJECT="${GOOGLE_CLOUD_PROJECT?}" \
         ZONE="${GOOGLE_CLOUD_ZONE?}" \
         CLUSTER_NAME="${GOOGLE_CLOUD_CLUSTER_NAME?}" \
         OUTPUT_DIR="${GOOGLE_CLOUD_BUCKET?}" \
         BASE_OUTPUT_DIR="${GOOGLE_CLOUD_BUCKET?}"'/output/' \
         DATASET_PATH="${GOOGLE_CLOUD_BUCKET?}"'/' \
         WORKLOAD='job_name_goes_here'

$ python3 -m MaxText.train MaxText/configs/base.yml \
      run_name="${USER}"'002' \
      base_output_directory="${OUTPUT_DIR?}" \
      dataset_type='synthetic' \
      steps='10' \
      model_name='llama2-7b'

$ command='python3 -m MaxText.train MaxText/configs/base.yml
      base_output_directory='"${BASE_OUTPUT_DIR?}"'
      dataset_path='"${DATASET_PATH?}"'
      steps=100
      per_device_batch_size=1'

$ xpk workload create \
      --base-docker-image 'maxtext_base_image' \
      --zone "${ZONE?}"
      --cluster "${CLUSTER_NAME?}" \
      --workload "${WORKLOAD?}" \
      --tpu-type='v6e-256' \
      --num-slices='1' \
      --command "${command?}"

$ python3 -m benchmarks.benchmark_runner xpk \
      --base_docker_image='maxtext_base_image' \
      --project="${PROJECT?}" \
      --zone="${ZONE?}" \
      --cluster_name="${CLUSTER_NAME?}" \
      --device_type='v6e-256' \
      --num_slices='1' \
      --base_output_directory="${OUTPUT_DIR?}" \
      --model_name="${MODEL_NAME?}"

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@SamuelMarks SamuelMarks force-pushed the setup-deps-refactor branch 3 times, most recently from 86b40cc to 7019d4b Compare October 24, 2025 03:18
@SamuelMarks SamuelMarks marked this pull request as ready for review October 24, 2025 03:20
@SamuelMarks SamuelMarks mentioned this pull request Oct 24, 2025
Merged
4 tasks
bvandermoon
bvandermoon reviewed Oct 28, 2025
View reviewed changes
@SamuelMarks SamuelMarks force-pushed the setup-deps-refactor branch 3 times, most recently from 03c95d4 to c8b8de8 Compare October 29, 2025 17:33
@kanglant
Copy link
Collaborator

kanglant commented Oct 29, 2025

Why do we need to relocate the Dockerfiles and the requirement files?

@SamuelMarks
Copy link
Collaborator Author

SamuelMarks commented Oct 29, 2025

Why do we need to relocate the Dockerfiles and the requirement files?

@kanglant This was agreed upon a long time ago in the RESTRUCTURE.md file. See internal design docs for more details.

@bvandermoon
Copy link
Collaborator

bvandermoon commented Oct 29, 2025

Why do we need to relocate the Dockerfiles and the requirement files?

@kanglant This was agreed upon a long time ago in the RESTRUCTURE.md file. See internal design docs for more details.

@SamuelMarks can we still have the base_requirements and generated_requirements structure within the new directory? These didn't exist at the time but it would help with organization to keep them

@SamuelMarks SamuelMarks force-pushed the setup-deps-refactor branch 6 times, most recently from 1cb7a0e to 0f56df0 Compare October 30, 2025 19:49
bvandermoon
bvandermoon reviewed Oct 30, 2025
View reviewed changes
@SamuelMarks SamuelMarks force-pushed the setup-deps-refactor branch 2 times, most recently from 5504888 to 42be91b Compare October 31, 2025 17:14
bvandermoon
bvandermoon approved these changes Oct 31, 2025
View reviewed changes
maxtext_runner.Dockerfile
# Copy assets separately
COPY src/MaxText/assets/ "${MAXTEXT_ASSETS_ROOT}"
COPY src/MaxText/test_assets/ "${MAXTEXT_TEST_ASSETS_ROOT}"
COPY generated_requirements .
Copy link
Collaborator

@bvandermoon bvandermoon Oct 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Was this removed intentionally? I guess it just isn't used?

Copy link
Collaborator

@shralex shralex Nov 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also would like to know. Thanks!

Copy link
Collaborator Author

@SamuelMarks SamuelMarks Nov 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That whole folder has been moved under dependencies so it's redundant and incorrect to retain.

shralex
shralex requested changes Nov 1, 2025
View reviewed changes
maxtext_runner.Dockerfile
# Copy assets separately
COPY src/MaxText/assets/ "${MAXTEXT_ASSETS_ROOT}"
COPY src/MaxText/test_assets/ "${MAXTEXT_TEST_ASSETS_ROOT}"
COPY generated_requirements .
Copy link
Collaborator

@shralex shralex Nov 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also would like to know. Thanks!

shralex
shralex reviewed Nov 1, 2025
View reviewed changes
Copy link
Collaborator

@shralex shralex left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added some more comments / questions.

Given the extensive changes to docker files and setup.sh, how are we testing these two ?

@bvandermoon bvandermoon mentioned this pull request Nov 3, 2025
Closed
4 tasks
@copybara-service copybara-service bot merged commit 83b3519 into AI-Hypercomputer:main Nov 3, 2025
65 checks passed
@bvandermoon bvandermoon mentioned this pull request Nov 4, 2025
Merged
4 tasks
gpupuck added a commit to NVIDIA/JAX-Toolbox that referenced this pull request Nov 4, 2025
@gpupuck
Just another casual day of refactoring: AI-Hypercomputer/maxtext#2541
17b8fc1
@gpupuck gpupuck mentioned this pull request Nov 4, 2025
Merged
gpupuck added a commit to NVIDIA/JAX-Toolbox that referenced this pull request Nov 4, 2025
@gpupuck
Just another casual day of refactoring MaxText PR#2541 (#1768)
8b3c50b 
AI-Hypercomputer/maxtext#2541
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@kanglant kanglant kanglant left review comments

@bvandermoon bvandermoon bvandermoon approved these changes

@shralex shralex shralex approved these changes

@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk is a code owner

@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner

@parambole parambole Awaiting requested review from parambole parambole is a code owner

@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner

@jacoguzo jacoguzo Awaiting requested review from jacoguzo

@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner

@gagika gagika Awaiting requested review from gagika gagika is a code owner

@A9isha A9isha Awaiting requested review from A9isha A9isha is a code owner

@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner

@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner

@mitalisi mitalisi Awaiting requested review from mitalisi mitalisi is a code owner

@gpolovets1 gpolovets1 Awaiting requested review from gpolovets1 gpolovets1 is a code owner

@mailvijayasingh mailvijayasingh Awaiting requested review from mailvijayasingh mailvijayasingh is a code owner

@jrplatin jrplatin Awaiting requested review from jrplatin jrplatin is a code owner

@patemotter patemotter Awaiting requested review from patemotter patemotter is a code owner

@Lumosis Lumosis Awaiting requested review from Lumosis Lumosis is a code owner

@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC is a code owner

@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo is a code owner

@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei is a code owner

@NuojCheng NuojCheng Awaiting requested review from NuojCheng NuojCheng is a code owner

@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 is a code owner

Assignees

No one assigned

Labels

pull ready

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@SamuelMarks @kanglant @bvandermoon @shralex